#include "mpi.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

#define N_BARR 0

struct comm_info 
{
    MPI_Comm       communicator;   /*   通信子           */
    MPI_Datatype   s_data_type;    /*   发送数据类型     */
    MPI_Datatype   r_data_type;    /*   接收数据类型     */
    int            rank;           /*   虚拟进程号       */
    void*          s_buffer;       /*   发送缓冲区       */
    size_t         s_alloc;        /*   发送缓冲区大小   */
    void*          r_buffer;       /*   接收缓冲区       */
    size_t         r_alloc;        /*   接收缓冲区大小   */
};

struct iter_schedule
{
    int      n_sample;       /*   采样次数   */
    int      s_cache_iter;   /*   可以忽略   */
    size_t   s_offs;         /*   可以忽略   */
};

// 散发延迟测试程序
// c_info:          MPI相关参数
// size  :          散发数据大小
// ITERATIONS:      测试采样次数
// time:            延迟 
void IMB_scatter(struct comm_info* c_info, int size, struct iter_schedule* ITERATIONS, double* time) 
{
    int ierr;
    double t1, t2;
    int    i,root;
    /*Type_Size s_size;*/
    int s_num, r_num;
    char *s_buf, *r_buf;

    ierr = 0;

    if(c_info->rank!=-1)
    {
        root=0;
        s_num = r_num = size;

        for(i=0; i<N_BARR; i++) 
            MPI_Barrier(c_info->communicator);

        t1 = MPI_Wtime();

        for(i=0; i< ITERATIONS->n_sample; i++)
        {
            /* 非根节点设定接收缓冲区，根节点设定发送缓冲区 */
            // 每个进程的rank是不一样的
            // 对于散发操作，发送缓冲区和接收缓冲区必须不同
            s_buf = c_info->s_buffer;
            r_buf = c_info->r_buffer;
            /* printf("rank is: %d\n", c_info->rank); */
            /* printf("root is: %d\n", root); */
            /* fflush(stdout); */

            MPI_Scatter(s_buf, s_num, c_info->s_data_type, r_buf, r_num, c_info->r_data_type, root, c_info->communicator);

#if defined(IMB_SCATTER_DBG)
            r_buf[s_num] = 0;
            // 该语句输出的次数与测试采样次数相同
            printf("[%d]: %s\n", c_info->rank, r_buf);
#endif
        }

        t2 = MPI_Wtime();
        *time=(t2 - t1)/(ITERATIONS->n_sample);
    }
    else
    {
        *time = 0.;
    }
}

// 广播延迟测试程序
// c_info:          MPI相关参数
// size  :          广播数据大小
// ITERATIONS:      测试采样次数
// time:            延迟 
void IMB_bcast(struct comm_info* c_info, int size, struct iter_schedule* ITERATIONS, double* time) 
{
    int ierr;
    double t1, t2;
    int    i,root;
    /*Type_Size s_size;*/
    int s_num;
    void* bc_buf;

    ierr = 0;

    if(c_info->rank!=-1)
    {
        root=0;
        s_num = size;
        for(i=0; i<N_BARR; i++) 
            MPI_Barrier(c_info->communicator);

        t1 = MPI_Wtime();
        /* 运行广播函数 n_sample 次 */
        for(i=0; i< ITERATIONS->n_sample; i++)
        {
            /* 非根节点设定接收缓冲区，根节点设定发送缓冲区 */
            // 每个进程的rank是不一样的
            bc_buf = (root == c_info->rank) ? c_info->s_buffer : c_info->r_buffer;
            /* printf("rank is: %d\n", c_info->rank); */
            /* printf("root is: %d\n", root); */
            /* fflush(stdout); */

            ierr= MPI_Bcast((char*)bc_buf,
                                    s_num,                   
                                    c_info->s_data_type,
                                    root,
                                    c_info->communicator);
#if defined(IMB_BCAST_DBG_EAGER) || defined(IMB_DEBUG_EAGERSHORT)
            printf("[%d]: %s\n", c_info->rank, bc_buf);
#endif
        }

        t2 = MPI_Wtime();
        /* printf("t1 is: %lf\t t2 is: %lf\t t2-t1 is: %lf\n", t1, t2, t2-t1); */
        *time=(t2 - t1)/(ITERATIONS->n_sample);
    }
    else
    {
        *time = 0.;
    }
}

/* #define SND_BUFFER_SIZE 1024 */
/* #define RCV_BUFFER_SIZE 1024 */

void IMB_init(struct comm_info* c_info, int size, struct iter_schedule* ITERATIONS, int msg_sz) 
{
    /*c_info->rank = 0;*/
    c_info->communicator = MPI_COMM_WORLD;
    c_info->s_data_type = c_info->r_data_type = MPI_CHAR;
    c_info->s_alloc = c_info->r_alloc = msg_sz+4;

    // 为发送缓冲区分配内存
    c_info->s_buffer = malloc(c_info->s_alloc);
#ifdef IMB_BCAST_DBG_EAGERSHORT
    memcpy(c_info->s_buffer, "hello world", strlen("hello world")+1);
#endif
#ifdef IMB_BCAST_DBG_EAGER
    memcpy(c_info->s_buffer, "too young too simple, sometimes naive", strlen("too young too simple, sometimes naive")+1);
#endif
#ifdef IMB_SCATTER_DBG
    char *msg = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
    memcpy(c_info->s_buffer, msg, strlen(msg)+1);
#endif
    // 为接收缓冲区分配内存
    c_info->r_buffer = malloc(c_info->r_alloc);
}

void IMB_finalize(struct comm_info* c_info, int size, struct iter_schedule* ITERATIONS, int msg_sz) 
{
    // free memory alloced on heap
}

#define TEST_UNDEFINED -1
#define TEST_BCAST_LATENCY 0
#define TEST_SCATTER_LATENCY 1

int main(int argc, char *argv[])
{
    int proc_nums;
    int my_rank;

    struct comm_info C_INFO;
    struct iter_schedule iter_sdu;
    int msg_sz;
    int sample_times;
    double time;
    int ops;

    if (argc != 4)
    {
        printf("Usage: %s MsgSize SampleTimes Operations\n", argv[0]);
        exit(-1);
    }
    else 
    {
        msg_sz = atoi(argv[1]);
        sample_times = atoi(argv[2]);
        if (strcmp(argv[3], "BCAST_Latency") == 0) {
            ops = TEST_BCAST_LATENCY;  
        } else if (strcmp(argv[3], "SCATTER_Latency") == 0) {
            ops = TEST_SCATTER_LATENCY;
        } else {
			ops = TEST_UNDEFINED;	
		}
			
    }

    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &proc_nums);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);
    
    /* 测试的采样次数 */
    iter_sdu.n_sample = sample_times;
    iter_sdu.s_cache_iter = iter_sdu.s_offs = 0;
    C_INFO.rank = my_rank;

    IMB_init(&C_INFO, 0, &iter_sdu, msg_sz);
    // 不要进行过barrier操作，否则会导致程序无法正常终止，索性影响不大
    // 服务器测试程序时还是需要barrier的
    MPI_Barrier(MPI_COMM_WORLD);

    switch (ops) {
        case TEST_BCAST_LATENCY:
            /* printf("BCAST_Latency: msg_sz: %d, sample_times: %d\n", msg_sz, sample_times); */
            IMB_bcast(&C_INFO, msg_sz, &iter_sdu, &time);
			printf("BCAST:P[%d]:\tMsg size: %d\t Latency: %lf\n", my_rank, msg_sz, time);
            break;
        case TEST_SCATTER_LATENCY:
            /* printf("SCATTER_Latency: msg_sz: %d, sample_times: %d\n", msg_sz, sample_times); */
            IMB_scatter(&C_INFO, msg_sz, &iter_sdu, &time);
			printf("SCATTER:P[%d]:\tMsg size: %d\t Latency: %lf\n", my_rank, msg_sz, time);
            break;
		case TEST_UNDEFINED:
        default:
			printf("Test undefined!\n");
			exit(-1);
            break;            
    }

    // 输出测试结果

    IMB_finalize(&C_INFO, 0, &iter_sdu, msg_sz);

    // 注意，目前不能调用该操作，因为里面会用到bcast操作，可能会失败
    // MPI_Finalize();

    return 0;
}
